// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/file/file_systems.h>
#include <pollux/core/plan_node.h>
#include <pollux/core/query_ctx.h>
#include <pollux/parse/plan_nodeId_generator.h>

namespace kumo::pollux::exec {
    class Task;
}

namespace kumo::pollux::tool::trace {
    class OperatorReplayerBase {
    public:
        OperatorReplayerBase(
                std::string traceDir,
                std::string queryId,
                std::string taskId,
                std::string nodeId,
                std::string operatorType,
                const std::string &driverIds,
                uint64_t queryCapacity,
                melon::Executor *executor);

        virtual ~OperatorReplayerBase() = default;

        OperatorReplayerBase(const OperatorReplayerBase &other) = delete;

        OperatorReplayerBase &operator=(const OperatorReplayerBase &other) = delete;

        OperatorReplayerBase(OperatorReplayerBase &&other) noexcept = delete;

        OperatorReplayerBase &operator=(OperatorReplayerBase &&other) noexcept =
        delete;

        virtual RowVectorPtr run(bool copyResults = true);

    protected:
        virtual core::PlanNodePtr createPlanNode(
                const core::PlanNode *node,
                const core::PlanNodeId &nodeId,
                const core::PlanNodePtr &source) const = 0;

        core::PlanNodePtr createPlan();

        std::shared_ptr<core::QueryCtx> createQueryCtx();

        const std::string queryId_;
        const std::string taskId_;
        const std::string nodeId_;
        const std::string operatorType_;
        const std::string taskTraceDir_;
        const std::string nodeTraceDir_;
        const std::shared_ptr<filesystems::FileSystem> fs_;
        const std::vector<uint32_t> pipelineIds_;
        const std::vector<uint32_t> driverIds_;
        const uint64_t queryCapacity_;
        const std::shared_ptr<core::PlanNodeIdGenerator> planNodeIdGenerator_{
                std::make_shared<core::PlanNodeIdGenerator>()};
        melon::Executor *const executor_;

        std::unordered_map<std::string, std::string> queryConfigs_;
        std::unordered_map<std::string, std::unordered_map<std::string, std::string>>
                connectorConfigs_;
        core::PlanNodePtr planFragment_;
        core::PlanNodeId replayPlanNodeId_;

        void printStats(const std::shared_ptr<exec::Task> &task) const;

    private:
        std::function<core::PlanNodePtr(std::string, core::PlanNodePtr)>
        replayNodeFactory(const core::PlanNode *node) const;
    };
} // namespace kumo::pollux::tool::trace
